#!/usr/bin/python3
# Written by Zen Shawn

from threading import Thread
from typing import Dict, List
import pickle
import time

import numpy as np
import vtk
import socket
import select

from .XmlParser import RobotTree
from .modern_robotics import FKinSpace, IKinSpace
from .ZenLogger import ZenLogger
from .Trajectory import ComputeTrajectory

logger = ZenLogger.get_logger(__name__, ZenLogger.INFO)


class Controller:
    def __init__(self, robot_actors: Dict[str, vtk.vtkActor], robot_tree: RobotTree, **kwargs):
        self.motion_server = None
        self.robot_actors = robot_actors
        self.robot_tree = robot_tree
        self.dof = self.robot_tree.GetDoF()
        self.interp_period = 0.050
        self.q_cur = [0.0 for _ in range(self.dof)]
        self.link_init_poses = {}
        links = self.robot_tree.m_links
        joints = self.robot_tree.m_joints
        for node in robot_tree.GetChain():
            link = links.get(node, False)
            if not link or link.get("visual_mesh", "") == "":
                continue
            if link["parent"]:
                parent_joint_pos = joints[link["parent"]]["origin_xyz"]
                parent_joint_rpy = joints[link["parent"]]["origin_rpy"]
                # link_to_parent = np.eye(4)
                # link_to_parent[:3, 3] = parent_joint_pos
                # link_to_parent[:3, :3] = Rotation.from_euler(
                #     "zyx", parent_joint_rpy).as_matrix()
                # link_to_base = link_to_base @ link_to_parent
                transform = vtk.vtkTransform()
                transform.Translate(parent_joint_pos)
                transform.RotateZ(parent_joint_rpy[2])
                transform.RotateX(parent_joint_rpy[1])
                transform.RotateY(parent_joint_rpy[0])
                self.link_init_poses[link['name']] = transform

        jls_lu = np.ones((self.robot_tree.GetDoF(), 2))
        jls = np.ones((3, self.robot_tree.GetDoF()))  # vel, acc, jerk
        for i, joint_name in enumerate(self.robot_tree.active_joints):
            jls[0, i] = joints[joint_name]['limit'].get('velocity', 1.0)
            jls[1, i] = joints[joint_name]['limit'].get('acceleration', 1.0)
            jls[2, i] = joints[joint_name]['limit'].get('jerk', 1.0)
            jls_lu[i, 0] = joints[joint_name]['limit'].get('lower', -1.57)
            jls_lu[i, 1] = joints[joint_name]['limit'].get('upper', 1.57)
        self.joint_limits = jls
        self.vlim = np.c_[-self.joint_limits[0], self.joint_limits[0]].T
        self.vlim = np.rad2deg(self.vlim)
        self.alim = np.c_[-self.joint_limits[1], self.joint_limits[1]].T
        self.alim = np.rad2deg(self.alim)
        self.joint_position_limits = jls_lu
        self.target_joint = [0.0 for _ in range(self.dof)]
        self.target_joint_updated = False

        self.traj_start_time = time.time()
        self.traj = None

        self.server = None

    def __del__(self):
        self.StopServer()

    def __ServoJ(self, q):
        logger.debug(f"Servo to: {q}")
        self.target_joint = q
        self.target_joint_updated = True

    def GetJointPosition(self):
        return self.q_cur

    def StartServer(self):
        # A UDP server
        if self.server is None:
            self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.server.bind(("0.0.0.0", 52416))
            self.server.setblocking(False)
            # self.thrd = Thread(target=self._Update, daemon=False)
            # self.thrd.start()

    def _Sync(self):
        if not self.target_joint_updated:
            return
        logger.debug(f"sync to : {self.target_joint}")
        for i, (axis, link_name) in enumerate(zip(self.robot_tree.active_joints_axis, self.robot_tree.active_links)):
            actor = self.robot_actors[link_name]
            transform: vtk.vtkTransform = actor.GetUserTransform()
            transform.SetMatrix(self.link_init_poses[link_name].GetMatrix())
            transform.RotateWXYZ(self.target_joint[i], *axis)
            self.q_cur[i] = self.target_joint[i]
        self.target_joint_updated = False

    def _Update(self):
        self._ProcessRequst()
        # self._FollowTraj()
        self._Sync()

    def _ProcessRequst(self):
        logger.debug(f"Waiting client request...")
        if self.server is None:
            return
        server = self.server
        try:
            data, addr = server.recvfrom(1024)
        except BlockingIOError:
            logger.debug("There is no data")
            return
        request = pickle.loads(data)
        msg_type = request.get("MsgType", "")
        logger.info(f"Received request: {request}")
        if msg_type == "MoveJ":
            q = request.get("Target", [])
            if len(q) == self.dof:
                self.MoveJ(q)
                server.sendto(b"OK", 0, addr)
                return
            else:
                server.sendto(
                    f"DoF({len(q)}) of requested data != DoF({self.dof}) of controller".encode(), 0, addr)
        elif msg_type == "ServoJ":
            q = request.get("Target", [])
            if len(q) == self.dof:
                self.__ServoJ(q)
                server.sendto(b"OK", 0, addr)
                return
            else:
                server.sendto(
                    f"DoF({len(q)}) of requested data != DoF({self.dof}) of controller".encode(), 0, addr)
        elif msg_type == "GetPosition":
            q = self.GetJointPosition()
            data = {"JointPosition": list(q)}
            server.sendto(pickle.dumps(data), 0, addr)
        else:
            server.sendto(
                f"Unknown message type {msg_type}".encode("utf8"), 0, addr)

    def StopServer(self):
        logger.debug(f"StopServer...")
        if self.server is not None:
            self.server.close()
            self.server = None

    def RandomState(self, wait=False):
        joints = np.random.uniform(
            self.joint_position_limits[:, 0], self.joint_position_limits[:, 1], self.robot_tree.GetDoF())
        return self.MoveJ(np.rad2deg(joints), wait)

    def FK(self, joints):
        pose = FKinSpace(self.robot_tree.m_home_pose,
                         self.robot_tree.m_s_list, joints)
        return pose

    def IK(self, pose):
        joints = IKinSpace(self.robot_tree.m_s_list, self.robot_tree.m_home_pose,
                           pose, self.GetJointPosition(), 0.01, 0.001)
        return joints

    def MoveJ(self, joints, wait=False):
        """Move robot joints

        Arguments:
            joints {list, np.array} -- target joints, dimension: dof

        Returns:
            bool -- result
        """
        traj = ComputeTrajectory(
            np.array([self.GetJointPosition(), joints]), self.vlim, self.alim)
        # self.traj = traj
        # self.traj_start_time = time.time()
        return self.FollowTraj(traj, wait)

    def FollowTraj(self, traj, wait=False):
        """Follow a toppra trajecotry

        Arguments:
            traj {class} -- joint trajectory, should have property "duration" and "eval" method
            wait (bool, optional): blocking or new thread. Defaults to False.

        Returns:
            bool -- result
        """
        if not wait:
            thread = Thread(target=self.FollowTraj, args=(
                traj,), kwargs={"wait": True}, daemon=False)
            thread.start()
            return True
        t0 = time.time()
        t = 0
        while t < traj.duration:
            self.__ServoJ(traj.eval(t))
            t = time.time() - t0
            time.sleep(self.interp_period)
        return True

    def _FollowTraj(self):
        if self.traj is not None:
            t = time.time() - self.traj_start_time
            if t >= self.traj.duration:
                self.__ServoJ(self.traj.eval(self.traj.duration))
                self.traj = None
            else:
                self.__ServoJ(self.traj.eval(t))


if __name__ == "__main__":
    urdf_path = "/home/xiaodaxia/Rvbust/Sources/RobotZenVispy/models/GP12/GP12.urdf"
    robot_tree = RobotTree(urdf_path)
    controller = Controller(None, robot_tree)
    controller.StartServer()
    from IPython import embed
    embed()
